import os
import json
import random
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
import cv2
import glob
import scipy.io as io
import torchvision.transforms as standard_transforms
import warnings
warnings.filterwarnings('ignore')

class CC(Dataset):
    def __init__(self, data_root, transform=None, train=False, flip=False):
        self.root_path = data_root

        prefix = "train" if train else "test"
        self.prefix = prefix
        self.img_list = os.listdir(f"{data_root}/{prefix}/rgb")
        
        # get image and ground-truth list
        self.gt_list = {}
        for img_name in self.img_list:
            img_path = f"{data_root}/{prefix}/rgb/{img_name}"  
            gt_path = f"{data_root}/{prefix}/label/{img_name}"
            self.gt_list[img_path] = gt_path.replace("jpg", "json")
        
        self.img_list = sorted(list(self.gt_list.keys()))
        self.nSamples = len(self.img_list)
        
        self.transform_rgb = transform['rgb']
        self.transform_tir = transform['tir']
        self.train = train
        self.flip = flip
        self.patch_size = 256
    
    def compute_density(self, points):
        """
        Compute crowd density:
            - defined as the average nearest distance between ground-truth points
        """
        points_tensor = torch.from_numpy(points.copy())
        dist = torch.cdist(points_tensor, points_tensor, p=2)
        if points_tensor.shape[0] > 1:
            density = dist.sort(dim=1)[0][:,1].mean().reshape(-1)
        else:
            density = torch.tensor(999.0).reshape(-1)
        return density

    def __len__(self):
        return self.nSamples

    def __getitem__(self, index):
        assert index <= len(self), 'index range error'

        # load image and gt points
        img_path = self.img_list[index]
        gt_path = self.gt_list[img_path]
        img, tir, points = load_data((img_path, gt_path), self.train)
        points = points.astype(float)

        # image transform
        if self.transform_rgb is not None:
            img = self.transform_rgb(img)
            tir = self.transform_tir(tir)
        img = torch.Tensor(img)
        tir = torch.Tensor(tir)

        # random scale
        """
        if self.train:
            scale_range = [0.8, 1.2]           
            min_size = min(img.shape[1:])
            scale = random.uniform(*scale_range)
            
            # interpolation
            if scale * min_size > self.patch_size:  
                img = torch.nn.functional.upsample_bilinear(img.unsqueeze(0), scale_factor=scale).squeeze(0)
                tir = torch.nn.functional.upsample_bilinear(tir.unsqueeze(0), scale_factor=scale).squeeze(0)
                points *= scale
        """

        # random crop patch
        if self.train:
            img, tir, points = random_crop(img, tir, points, patch_size=self.patch_size)

        # random flip
        """
        if random.random() > 0.5 and self.train and self.flip:
            img = torch.flip(img, dims=[2])
            tir = torch.flip(tir, dims=[2])
            points[:, 1] = self.patch_size - points[:, 1]
        """
        
        # target
        target = {}
        target['points'] = torch.Tensor(points)
        target['labels'] = torch.ones([points.shape[0]]).long()

        if self.train:
            density = self.compute_density(points)
            target['density'] = density

        if not self.train:
            target['image_path'] = img_path

        return (img, tir), target


def load_data(img_gt_path, train):
    img_path, gt_path = img_gt_path
    tir_path = img_path.replace('/rgb/', '/tir/')
    img = cv2.imread(img_path)
    tir = cv2.imread(tir_path)
    img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    tir = Image.fromarray(cv2.cvtColor(tir, cv2.COLOR_BGR2RGB))
    
    # points = io.loadmat(gt_path)['image_info'][0][0][0][0][0][:,::-1]
    
    with open(gt_path, 'r') as file:
        lines = [line for line in file]
        if len(lines) > 1:
            print(f"[ERROR]load_data: len(lines)={len(lines)}, file={gt_path}")
            exit
        label = json.loads(lines[0])
    points = np.array(label['points'])
    return img, tir, points


def random_crop(img, tir, points, patch_size=256):
    patch_h = patch_size
    patch_w = patch_size
    
    # random crop
    start_h = random.randint(0, img.size(1) - patch_h) if img.size(1) > patch_h else 0
    start_w = random.randint(0, img.size(2) - patch_w) if img.size(2) > patch_w else 0
    end_h = start_h + patch_h
    end_w = start_w + patch_w
    idx = (points[:, 0] >= start_h) & (points[:, 0] <= end_h) & (points[:, 1] >= start_w) & (points[:, 1] <= end_w)

    # clip image and points
    result_img = img[:, start_h:end_h, start_w:end_w]
    result_tir = tir[:, start_h:end_h, start_w:end_w]
    result_points = points[idx]
    result_points[:, 0] -= start_h
    result_points[:, 1] -= start_w
    
    # resize to patchsize
    """
    imgH, imgW = result_img.shape[-2:]
    fH, fW = patch_h/imgH, patch_w/imgW
    result_img = torch.nn.functional.interpolate(result_img.unsqueeze(0), (patch_h, patch_w)).squeeze(0)
    result_tir = torch.nn.functional.interpolate(result_tir.unsqueeze(0), (patch_h, patch_w)).squeeze(0)
    result_points[:, 0] *= fH
    result_points[:, 1] *= fW
    """
    return result_img, result_tir, result_points


def build(image_set, args):
    transform = {
        'rgb': standard_transforms.Compose([
                   standard_transforms.ToTensor(),
                   standard_transforms.Normalize(mean=[0.407, 0.389, 0.396], std=[0.241, 0.246, 0.242]),]),
        'tir': standard_transforms.Compose([
                   standard_transforms.ToTensor(),
                   standard_transforms.Normalize(mean=[0.492, 0.168, 0.430], std=[0.317, 0.174, 0.191]),])
    }
    
    data_root = args.data_path
    if image_set == 'train':
        train_set = CC(data_root, train=True, transform=transform, flip=True)
        return train_set
    elif image_set == 'val':
        val_set = CC(data_root, train=False, transform=transform)
        return val_set
